#!/usr/local/bin/python
#-*- coding: utf-8 -*-
#############################################################################
#                                                                           #
#   File: proxy.py                                                          #
#                                                                           #
#   Copyright (C) 2008-2011 lyricconch <lyricconch@gmail.com>               #
#                                                                           #
#   This file is part of WallProxyMod.                                      #
#                                                                           #
#   WallProxyMod is free software: you can redistribute it and/or modify    #
#   it under the terms of the GNU General Public License as                 #
#   published by the Free Software Foundation, either version 3 of the      #
#   License, or (at your option) any later version.                         #
#                                                                           #
#   WallProxyMod is distributed in the hope that it will be useful,         #
#   but WITHOUT ANY WARRANTY; without even the implied warranty of          #
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the           #
#   GNU General Public License for more details.                            #
#                                                                           #
#   You should have received a copy of the GNU General Public License       #
#   along with WallProxyMod.  If not, see <http://www.gnu.org/licenses/>.   #
#                                                                           #
#############################################################################

__author__ = "lyricconch@gmail.com"
__version__ = "2.0.1a"

from multitask import taskmgr, recv, accept, sleep
from util import recvall, sendall, connect, gae_fetch

import socket
import urlparse, re

def requires_proxy(addr, method, path, head):
    return True

def client_handler(sock, addr):
    prefix = yield recvall(sock, "\r\n\r\n")
    line, CRLF, head = prefix.partition("\r\n")
    print "%s:%s want " % addr + line.strip()
    method, path, ver = line.split()
    if requires_proxy(addr, method, path, head) and method != "CONNECT":
        body = yield recvall(sock, get_length(head))
        ret = yield gae_fetch(prefix + body)
        yield sendall(sock, ret)
        sock.shutdown(socket.SHUT_WR)
        return
    elif method == "CONNECT":
        remote = yield connect(*path.rsplit(":", 1))
        yield sendall(sock, "HTTP/1.0\r\n 200 Connection Made\r\n\r\n")
        taskmgr.add(pass_through(sock, remote))
        taskmgr.add(pass_through(remote, sock))
        return
    else:
        rs = urlparse.urlsplit(path)
        if rs.scheme != "http" or not rs.netloc: return
        path = urlparse.urlunsplit(("", "", rs.path, rs.query, ""))
        head = _re_conn.sub("" , head)
        format = "%s %s %s\r\nConnection: close\r\n%s"
        remote = yield connect(rs.hostname, rs.port or 80)
        yield sendall(remote, format % (method, path, ver, head))
        taskmgr.add(pass_through(sock, remote, get_length(head)))
        yield http_session(remote, sock)
        # continue for pipelining?
    remote.shutdown(socket.SHUT_RD)
    sock.shutdown(socket.SHUT_WR)

def http_session(source, dest):
    prefix = yield recvall(source, "\r\n\r\n")
    yield sendall(dest, prefix)
    line, CRLF, head = prefix.partition("\r\n")
    length = get_length(head)
    if length is not None:
        yield pass_through(source, dest, length)
    else:
        yield chunked_transfer(source, dest)

_re_conn = re.compile(r"(Proxy-)?Connection:.*\n?", re.I)
_re_ctnlen = re.compile(r"Content-Length:\s*(\d+)", re.I)
_re_transenc = re.compile(r"Transfer-Encoding:\s*(\S+)", re.I)
def get_length(head):
    chunked = _re_transenc.findall(head)
    if chunked:
        if chunked[0] == "chunked": return None
        raise RuntimeError("unknown transfer-encoding" + chunked[0])
    length = _re_ctnlen.findall(head)
    return int(length[0]) if length else 0

def pass_through(local, remote, length= -1):
    while length != 0:
        data = yield recv(local, length if 0 < length < 16384 else 16384)
        if not data:
            if length > 0: print "bad request: not finished body"
            return
        length -= len(data)
        yield sendall(remote, data)

def chunked_transfer(source, dest, transform=False):
    while True:
        line = yield recvall(source, "\r\n")
        size = int(line.split(";")[0], 16)
        if size == 0: break
        if not transform: yield sendall(dest, line)
        data = yield recvall(source, size + 2)
        yield sendall(dest, data[:-2] if transform else data)
    while not line.endswith("\r\n\r\n"):
        line += yield recvall(source, "\r\n")
    if not transform: yield sendall(source, line)

def server_handler(address):
    sock = socket.socket()
    sock.bind(address)
    sock.listen(5)
    while True:
        conn, addr = yield accept(sock)
        taskmgr.add(client_handler(conn, addr))

def ping():
    while True:
        yield sleep(1)
        print "z",    
taskmgr.add(ping())

LISTEN_ADDRESS = ("", 8086)
taskmgr.add(server_handler(LISTEN_ADDRESS))
taskmgr.run()
